#!/usr/bin/env eco

local time = require 'eco.time'

local work_queue = {}
work_queue.__index = work_queue

-- 创建新的工作队列实例
function work_queue.new()
    local self = setmetatable({}, work_queue)
    self.tasks = {}        -- 普通任务队列
    self.timed_tasks = {}  -- 存储定时队列
    self.running = false
    self.workers = {}
    self.next_task_id = 1  -- 唯一任务ID
    self.task_handles = {} -- 句柄
    return self
end

-- 添加普通任务，返回任务句柄
function work_queue:add_task(func, ...)
    local args = {...}
    local task_id = self.next_task_id
    self.next_task_id = self.next_task_id + 1
    local task = {
        id = task_id,
        func = func,
        args = args,
        type = "normal"
    }
    table.insert(self.tasks, task)
    local handle = {
        id = task_id,
        queue = self,
        canceled = false
    }
    self.task_handles[task_id] = handle
    return handle
end

-- 添加紧急任务，返回任务句柄
function work_queue:add_urgent_task(func, ...)
    local args = {...}
    local task_id = self.next_task_id
    self.next_task_id = self.next_task_id + 1
    local task = {
        id = task_id,
        func = func,
        args = args,
        type = "urgent"
    }
    table.insert(self.tasks, 1, task)
    local handle = {
        id = task_id,
        queue = self,
        canceled = false
    }
    self.task_handles[task_id] = handle
    return handle
end

-- 添加定时任务，返回任务句柄
function work_queue:add_task_time(func, delay, ...)
    local args = {...}
    local task_id = self.next_task_id
    self.next_task_id = self.next_task_id + 1
    local task = {
        id = task_id,
        func = func,
        args = args,
        type = "timed",
        execute_time = os.time() + delay
    }
    table.insert(self.timed_tasks, task)
    table.sort(self.timed_tasks, function(a, b) 
        return a.execute_time < b.execute_time 
    end)
    local handle = {
        id = task_id,
        queue = self,
        canceled = false
    }
    self.task_handles[task_id] = handle
    return handle
end

-- 取消任务
function work_queue:cancel_task(handle)
    if handle and not handle.canceled then
        handle.canceled = true
        for i, task in ipairs(self.tasks) do
            if task.id == handle.id then
                table.remove(self.tasks, i)
                self.task_handles[handle.id] = nil
                return true
            end
        end
        for i, task in ipairs(self.timed_tasks) do
            if task.id == handle.id then
                table.remove(self.timed_tasks, i)
                self.task_handles[handle.id] = nil
                return true
            end
        end
    end
    return false
end

-- 检查定时任务是否需要执行
function work_queue:check_timed_tasks()
    local current_time = os.time()
    local moved_count = 0
    for i = #self.timed_tasks, 1, -1 do
        local task = self.timed_tasks[i]
        local handle = self.task_handles[task.id]
        if handle and handle.canceled then
            table.remove(self.timed_tasks, i)
            self.task_handles[task.id] = nil
        -- 如果任务到达执行时间，移到普通队列
        elseif task.execute_time <= current_time then
            table.remove(self.timed_tasks, i)
            task.type = "normal"
            table.insert(self.tasks, task)
            moved_count = moved_count + 1
        end
    end
    return moved_count
end

-- 工作线程函数
function work_queue:worker(id)
    while self.running do
        self:check_timed_tasks()
        if #self.tasks > 0 then
            local task = table.remove(self.tasks, 1)
            local handle = self.task_handles[task.id]
            if not handle or not handle.canceled then
                -- print(string.format("Worker %d processing task %d", id, task.id))
                if #task.args > 0 then
                    task.func(table.unpack(task.args))
                else
                    task.func()
                end
                -- 任务完成后移除句柄
                self.task_handles[task.id] = nil
            else
                -- 任务已被取消，移除句柄
                self.task_handles[task.id] = nil
            end
        else
            time.sleep(0.1)
        end
    end
end

-- 启动工作队列
function work_queue:start(num_workers)
    self.running = true
    for i = 1, num_workers do
        local worker = eco.run(function()
            self:worker(i)
        end)
        table.insert(self.workers, worker)
    end
end

-- 停止工作队列
function work_queue:stop()
    self.running = false
end

-- 获取队列状态
function work_queue:get_status()
    return {
        normal_tasks = #self.tasks,
        timed_tasks = #self.timed_tasks,
        active_handles = self.next_task_id - 1,
        running = self.running
    }
end

-- 创建默认系统队列
local sys_queue = work_queue.new()
sys_queue:start(8)

-- 返回模块
local M = {
    new = work_queue.new,  -- 允许用户创建新的工作队列
    sys_queue = sys_queue  -- 提供默认的系统队列
}

return M